import torch
from data import mydataloader
from models.StyleGAN import StyleGAN
from models.DCGAN import Discriminator

for i_batch, batch_data in enumerate(mydataloader):
    print(i_batch)
    print(batch_data.shape)
    b,c,h,w = batch_data.shape
    latent_dim=64

    stylegan = StyleGAN(64)
    discriminator = Discriminator()
    loss_func = torch.nn.MSELoss()
    gen_img = stylegan(torch.randn((b,64)))
    print(gen_img.shape)
    predict = discriminator(gen_img)
    print(predict.shape)
    loss = loss_func(gen_img,torch.zeros_like(gen_img))
    print(loss)
    loss.backward()
    print('done!!!')
    break
